/*
 * Copyright 2023-2025, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 *
 * 龙芯LoongArch64架构TLB管理实现
 * 本文件实现LoongArch64架构的TLB相关功能，包括TLB刷新、TLB条目操作等
 */

#include <KernelExport.h>

#include <arch/cpu.h>
#include <arch/loongarch64/arch_mmu.h>
#include <boot/kernel_args.h>
#include <kernel.h>
#include <vm/vm.h>
#include <vm/VMAddressSpace.h>

#include "LoongArchVMTranslationMap.h"

// 定义调试输出宏
#define TRACE_TLB
#ifdef TRACE_TLB
#define TRACE(x...) dprintf("arch_mmu_tlb: " x)
#else
#define TRACE(x...) ;
#endif

// 定义LoongArch64 TLB相关CSR寄存器
#define CSR_TLBIDX    0x10    // TLB索引
#define CSR_TLBEHI    0x11    // TLB入口高位
#define CSR_TLBELO0   0x12    // TLB入口低位0
#define CSR_TLBELO1   0x13    // TLB入口低位1
#define CSR_ASID      0x18    // 地址空间标识符

// 定义TLB操作指令
#define TLBSRCH     0x0       // 搜索TLB
#define TLBRD       0x1       // 读取TLB
#define TLBWR       0x2       // 写入TLB
#define TLBFILL     0x3       // 填充TLB
#define TLBCLR      0x4       // 清除TLB
#define TLBFLUSH    0x5       // 刷新TLB
#define INVTLB      0x6       // 无效化TLB

// 定义TLBIDX寄存器位
#define TLBIDX_PS_SHIFT  24    // 页大小偏移
#define TLBIDX_PS_MASK   0x3f  // 页大小掩码
#define TLBIDX_INDEX_SHIFT 0   // 索引偏移
#define TLBIDX_INDEX_MASK 0xfff // 索引掩码

// 定义INVTLB操作码
#define INVTLB_ALL          0    // 无效化所有TLB
#define INVTLB_ASID         1    // 无效化指定ASID的TLB
#define INVTLB_VA           2    // 无效化指定VA的TLB
#define INVTLB_ASID_VA      3    // 无效化指定ASID和VA的TLB
#define INVTLB_ASID_VA_PAGE 4    // 无效化指定ASID和VA页的TLB

// 定义页大小
#define PS_4K   12   // 4KB页

// 全局变量
static uint32 sTlbEntries = 0;  // TLB条目数量
static uint32 sTlbPageSizes = 0; // 支持的页大小

/**
 * 初始化TLB
 * 检测TLB条目数量和支持的页大小
 */
status_t
arch_mmu_tlb_init(kernel_args *args)
{
    TRACE("Initializing TLB\n");

    // 检测TLB条目数量
    // 通过读取CPUCFG寄存器获取TLB信息
    uint32 cpucfg2 = read_cpucfg(2);
    sTlbEntries = ((cpucfg2 >> 16) & 0x3ff) + 1;
    
    // 检测支持的页大小
    uint32 cpucfg3 = read_cpucfg(3);
    sTlbPageSizes = (cpucfg3 >> 8) & 0xff;

    TRACE("TLB entries: %u, supported page sizes: 0x%x\n", sTlbEntries, sTlbPageSizes);

    return B_OK;
}

/**
 * 初始化每个CPU的TLB
 * 清除TLB并设置初始状态
 */
status_t
arch_mmu_tlb_init_percpu(kernel_args *args, int cpu)
{
    TRACE("Initializing TLB for CPU %d\n", cpu);

    // 清除TLB
    asm volatile("invtlb %0, $zero, $zero" : : "i"(INVTLB_ALL) : "memory");

    return B_OK;
}

/**
 * 刷新整个TLB
 * 无效化所有TLB条目
 */
void
arch_mmu_tlb_flush(void)
{
    TRACE("Flushing all TLB entries\n");
    asm volatile("invtlb %0, $zero, $zero" : : "i"(INVTLB_ALL) : "memory");
}

/**
 * 刷新指定ASID的TLB条目
 * @param asid 地址空间标识符
 */
void
arch_mmu_tlb_flush_asid(uint32 asid)
{
    TRACE("Flushing TLB entries for ASID %u\n", asid);
    asm volatile("invtlb %0, %1, $zero" : : "i"(INVTLB_ASID), "r"(asid) : "memory");
}

/**
 * 刷新指定虚拟地址的TLB条目
 * @param vaddr 虚拟地址
 */
void
arch_mmu_tlb_flush_va(addr_t vaddr)
{
    TRACE("Flushing TLB entry for VA %#" B_PRIxADDR "\n", vaddr);
    asm volatile("invtlb %0, $zero, %1" : : "i"(INVTLB_VA), "r"(vaddr) : "memory");
}

/**
 * 刷新指定ASID和虚拟地址的TLB条目
 * @param asid 地址空间标识符
 * @param vaddr 虚拟地址
 */
void
arch_mmu_tlb_flush_asid_va(uint32 asid, addr_t vaddr)
{
    TRACE("Flushing TLB entry for ASID %u, VA %#" B_PRIxADDR "\n", asid, vaddr);
    asm volatile("invtlb %0, %1, %2" : : "i"(INVTLB_ASID_VA), "r"(asid), "r"(vaddr) : "memory");
}

/**
 * 刷新指定ASID和虚拟地址页的TLB条目
 * @param asid 地址空间标识符
 * @param vaddr 虚拟地址
 */
void
arch_mmu_tlb_flush_asid_va_page(uint32 asid, addr_t vaddr)
{
    TRACE("Flushing TLB entry for ASID %u, VA page %#" B_PRIxADDR "\n", asid, vaddr);
    asm volatile("invtlb %0, %1, %2" : : "i"(INVTLB_ASID_VA_PAGE), "r"(asid), "r"(vaddr) : "memory");
}

/**
 * 读取TLB条目
 * @param index TLB索引
 * @param _hi 输出参数，TLB高位
 * @param _lo0 输出参数，TLB低位0
 * @param _lo1 输出参数，TLB低位1
 */
status_t
arch_mmu_tlb_read(uint32 index, uint64 *_hi, uint64 *_lo0, uint64 *_lo1)
{
    if (index >= sTlbEntries)
        return B_BAD_VALUE;

    // 设置TLB索引
    uint64 tlbidx = index & TLBIDX_INDEX_MASK;
    asm volatile("csrwr %0, %1" : : "r"(tlbidx), "i"(CSR_TLBIDX) : "memory");

    // 读取TLB条目
    asm volatile("tlbrd" : : : "memory");

    // 读取TLB寄存器值
    uint64 hi, lo0, lo1;
    asm volatile("csrrd %0, %1" : "=r"(hi) : "i"(CSR_TLBEHI));
    asm volatile("csrrd %0, %1" : "=r"(lo0) : "i"(CSR_TLBELO0));
    asm volatile("csrrd %0, %1" : "=r"(lo1) : "i"(CSR_TLBELO1));

    // 输出结果
    if (_hi) *_hi = hi;
    if (_lo0) *_lo0 = lo0;
    if (_lo1) *_lo1 = lo1;

    return B_OK;
}

/**
 * 写入TLB条目
 * @param index TLB索引
 * @param hi TLB高位
 * @param lo0 TLB低位0
 * @param lo1 TLB低位1
 */
status_t
arch_mmu_tlb_write(uint32 index, uint64 hi, uint64 lo0, uint64 lo1)
{
    if (index >= sTlbEntries)
        return B_BAD_VALUE;

    // 设置TLB索引
    uint64 tlbidx = index & TLBIDX_INDEX_MASK;
    tlbidx |= (uint64)PS_4K << TLBIDX_PS_SHIFT; // 设置页大小为4KB
    asm volatile("csrwr %0, %1" : : "r"(tlbidx), "i"(CSR_TLBIDX) : "memory");

    // 设置TLB寄存器值
    asm volatile("csrwr %0, %1" : : "r"(hi), "i"(CSR_TLBEHI) : "memory");
    asm volatile("csrwr %0, %1" : : "r"(lo0), "i"(CSR_TLBELO0) : "memory");
    asm volatile("csrwr %0, %1" : : "r"(lo1), "i"(CSR_TLBELO1) : "memory");

    // 写入TLB条目
    asm volatile("tlbwr" : : : "memory");

    return B_OK;
}

/**
 * 搜索TLB条目
 * @param vaddr 虚拟地址
 * @param asid 地址空间标识符
 * @param _index 输出参数，找到的TLB索引
 * @return 如果找到返回B_OK，否则返回B_ENTRY_NOT_FOUND
 */
status_t
arch_mmu_tlb_search(addr_t vaddr, uint32 asid, uint32 *_index)
{
    // 设置ASID
    uint64 oldAsid;
    asm volatile("csrrd %0, %1" : "=r"(oldAsid) : "i"(CSR_ASID));
    asm volatile("csrwr %0, %1" : : "r"((uint64)asid), "i"(CSR_ASID) : "memory");

    // 设置虚拟地址
    asm volatile("csrwr %0, %1" : : "r"(vaddr), "i"(CSR_TLBEHI) : "memory");

    // 搜索TLB
    asm volatile("tlbsrch" : : : "memory");

    // 恢复ASID
    asm volatile("csrwr %0, %1" : : "r"(oldAsid), "i"(CSR_ASID) : "memory");

    // 读取TLB索引
    uint64 tlbidx;
    asm volatile("csrrd %0, %1" : "=r"(tlbidx) : "i"(CSR_TLBIDX));

    // 检查是否找到
    if (tlbidx & (1ULL << 31)) {
        // 未找到，最高位为1
        return B_ENTRY_NOT_FOUND;
    }

    // 输出索引
    if (_index)
        *_index = tlbidx & TLBIDX_INDEX_MASK;

    return B_OK;
}

/**
 * 填充TLB条目
 * 使用TLBFILL指令填充TLB
 * @param vaddr 虚拟地址
 * @param lo0 TLB低位0
 * @param lo1 TLB低位1
 */
void
arch_mmu_tlb_fill(addr_t vaddr, uint64 lo0, uint64 lo1)
{
    // 设置页大小为4KB
    uint64 tlbidx = (uint64)PS_4K << TLBIDX_PS_SHIFT;
    asm volatile("csrwr %0, %1" : : "r"(tlbidx), "i"(CSR_TLBIDX) : "memory");

    // 设置虚拟地址
    asm volatile("csrwr %0, %1" : : "r"(vaddr), "i"(CSR_TLBEHI) : "memory");

    // 设置物理地址和属性
    asm volatile("csrwr %0, %1" : : "r"(lo0), "i"(CSR_TLBELO0) : "memory");
    asm volatile("csrwr %0, %1" : : "r"(lo1), "i"(CSR_TLBELO1) : "memory");

    // 填充TLB
    asm volatile("tlbfill" : : : "memory");
}

/**
 * 获取TLB条目数量
 * @return TLB条目数量
 */
uint32
arch_mmu_tlb_get_entries_count(void)
{
    return sTlbEntries;
}

/**
 * 获取支持的页大小
 * @return 支持的页大小位掩码
 */
uint32
arch_mmu_tlb_get_supported_page_sizes(void)
{
    return sTlbPageSizes;
}

/**
 * 处理TLB缺失异常
 * 当访问未映射的虚拟地址时调用
 * @param vaddr 触发异常的虚拟地址
 * @param isWrite 是否为写操作
 * @return 处理状态
 */
status_t
arch_mmu_tlb_miss_handler(addr_t vaddr, bool isWrite)
{
    TRACE("TLB miss at VA %#" B_PRIxADDR ", %s access\n", vaddr, isWrite ? "write" : "read");

    // 获取当前地址空间
    VMAddressSpace *addressSpace = VMAddressSpace::GetCurrent();
    if (!addressSpace) {
        panic("TLB miss with no current address space!\n");
        return B_ERROR;
    }

    // 获取转换映射
    LoongArchVMTranslationMap *map = (LoongArchVMTranslationMap *)addressSpace->TranslationMap();
    if (!map) {
        panic("TLB miss with no translation map!\n");
        return B_ERROR;
    }

    // 查询物理地址和标志
    phys_addr_t physicalAddress;
    uint32 flags;
    status_t status = map->Query(vaddr, &physicalAddress, &flags);
    if (status != B_OK) {
        // 地址未映射，交由页面错误处理器处理
        return B_ENTRY_NOT_FOUND;
    }

    // 检查写保护
    if (isWrite && (flags & B_WRITE_AREA) == 0) {
        // 写保护违规，交由页面错误处理器处理
        return B_PERMISSION_DENIED;
    }

    // 构建TLB条目
    uint64 lo0 = (physicalAddress & kPtePhysAddrMask);
    
    // 设置访问权限
    if ((flags & B_READ_AREA) != 0)
        lo0 |= kPteRead;
    if ((flags & B_WRITE_AREA) != 0 || (flags & B_KERNEL_WRITE_AREA) != 0)
        lo0 |= kPteWrite;
    if ((flags & B_EXECUTE_AREA) != 0)
        lo0 |= kPteExec;

    // 设置用户/内核访问权限
    if ((flags & B_USER_PROTECTION) != 0)
        lo0 |= kPteUser;

    // 设置全局标志（内核页面通常是全局的）
    if ((flags & B_KERNEL_AREA) != 0)
        lo0 |= kPteGlobal;

    // 设置有效位和访问/脏位
    lo0 |= kPteValid | kPteAccessed;
    if (isWrite)
        lo0 |= kPteDirty;

    // 填充TLB
    // 对于4KB页，我们使用相同的映射填充两个相邻的页表项
    uint64 lo1 = lo0;
    arch_mmu_tlb_fill(vaddr, lo0, lo1);

    return B_OK;
}